#include "dct.h"
#include "common.h"

static inline void pixel_sub_wxh(dctcoef *diff, int i_size,
                                 pixel *pix1, int i_pix1, pixel *pix2, int i_pix2)
{
    for (int y = 0; y < i_size; y++)
    {
        for (int x = 0; x < i_size; x++)
            diff[x + y * i_size] = pix1[x] - pix2[x];
        pix1 += i_pix1;
        pix2 += i_pix2;
    }
}

void sub4x4_dct(dctcoef dct[16], pixel *pix1, pixel *pix2)
{
    dctcoef d[16] = {0};
    dctcoef tmp[16] = {0};

    pixel_sub_wxh(d, 4, pix1, FENC_STRIDE, pix2, FDEC_STRIDE);

    // 进行快速蝶形计算
    for (int i = 0; i < 4; i++)
    {
        int s03 = d[i * 4 + 0] + d[i * 4 + 3];
        int s12 = d[i * 4 + 1] + d[i * 4 + 2];
        int d03 = d[i * 4 + 0] - d[i * 4 + 3];
        int d12 = d[i * 4 + 1] - d[i * 4 + 2];

        tmp[0 * 4 + i] = s03 + s12;
        tmp[1 * 4 + i] = 2 * d03 + d12;
        tmp[2 * 4 + i] = s03 - s12;
        tmp[3 * 4 + i] = d03 - 2 * d12;
    }

    for (int i = 0; i < 4; i++)
    {
        int s03 = tmp[i * 4 + 0] + tmp[i * 4 + 3];
        int s12 = tmp[i * 4 + 1] + tmp[i * 4 + 2];
        int d03 = tmp[i * 4 + 0] - tmp[i * 4 + 3];
        int d12 = tmp[i * 4 + 1] - tmp[i * 4 + 2];

        dct[i * 4 + 0] = s03 + s12;
        dct[i * 4 + 1] = 2 * d03 + d12;
        dct[i * 4 + 2] = s03 - s12;
        dct[i * 4 + 3] = d03 - 2 * d12;
    }
}

void sub8x8_dct(dctcoef dct[4][16], pixel *pix1, pixel *pix2)
{
    sub4x4_dct(dct[0], &pix1[0], &pix2[0]);
    sub4x4_dct(dct[1], &pix1[4], &pix2[4]);
    sub4x4_dct(dct[2], &pix1[4 * FENC_STRIDE + 0], &pix2[4 * FDEC_STRIDE + 0]);
    sub4x4_dct(dct[3], &pix1[4 * FENC_STRIDE + 4], &pix2[4 * FDEC_STRIDE + 4]);
}

void sub16x16_dct(dctcoef dct[16][16], pixel *pix1, pixel *pix2)
{
    sub8x8_dct(&dct[0], &pix1[0], &pix2[0]);
    sub8x8_dct(&dct[4], &pix1[8], &pix2[8]);
    sub8x8_dct(&dct[8], &pix1[8 * FENC_STRIDE + 0], &pix2[8 * FDEC_STRIDE + 0]);
    sub8x8_dct(&dct[12], &pix1[8 * FENC_STRIDE + 8], &pix2[8 * FDEC_STRIDE + 8]);
}

void zigzag_scan_4x4_frame(dctcoef level[16], dctcoef dct[16])
{
    level[0] = dct[0 * 4 + 0];
    level[1] = dct[1 * 4 + 0];
    level[2] = dct[0 * 4 + 1];
    level[3] = dct[0 * 4 + 2];
    level[4] = dct[1 * 4 + 1];
    level[5] = dct[2 * 4 + 0];
    level[6] = dct[3 * 4 + 0];
    level[7] = dct[2 * 4 + 1];
    level[8] = dct[1 * 4 + 2];
    level[9] = dct[0 * 4 + 3];
    level[10] = dct[1 * 4 + 3];
    level[11] = dct[2 * 4 + 2];
    level[12] = dct[3 * 4 + 1];
    level[13] = dct[3 * 4 + 2];
    level[14] = dct[2 * 4 + 3];
    level[15] = dct[3 * 4 + 3];
}

void zigzag_scan_2x2_dc(dctcoef level[4], dctcoef dct[4])
{
    level[0] = dct[0 * 2 + 0];
    level[1] = dct[1 * 2 + 0];
    level[2] = dct[0 * 2 + 1];
    level[3] = dct[1 * 2 + 1];
}

void zigzag_scan_2x4_dc(dctcoef level[8], dctcoef dct[8])
{
    level[0] = dct[0];
    level[1] = dct[2];
    level[2] = dct[1];
    level[3] = dct[4];
    level[4] = dct[6];
    level[5] = dct[3];
    level[6] = dct[5];
    level[7] = dct[7];
}

void idct_dequant_2x4_dconly(dctcoef dct[8], int dequant_mf[6][16], int i_qp)
{
    int a0 = dct[0] + dct[1];
    int a1 = dct[2] + dct[3];
    int a2 = dct[4] + dct[5];
    int a3 = dct[6] + dct[7];
    int a4 = dct[0] - dct[1];
    int a5 = dct[2] - dct[3];
    int a6 = dct[4] - dct[5];
    int a7 = dct[6] - dct[7];
    int b0 = a0 + a1;
    int b1 = a2 + a3;
    int b2 = a4 + a5;
    int b3 = a6 + a7;
    int b4 = a0 - a1;
    int b5 = a2 - a3;
    int b6 = a4 - a5;
    int b7 = a6 - a7;

    int dmf = dequant_mf[i_qp % 6][0] << i_qp / 6;
    dct[0] = ((b0 + b1) * dmf + 32) >> 6;
    dct[1] = ((b2 + b3) * dmf + 32) >> 6;
    dct[2] = ((b0 - b1) * dmf + 32) >> 6;
    dct[3] = ((b2 - b3) * dmf + 32) >> 6;
    dct[4] = ((b4 - b5) * dmf + 32) >> 6;
    dct[5] = ((b6 - b7) * dmf + 32) >> 6;
    dct[6] = ((b4 + b5) * dmf + 32) >> 6;
    dct[7] = ((b6 + b7) * dmf + 32) >> 6;
}

void idct_dequant_2x2_dconly(dctcoef dct[4], int dequant_mf[6][16], int i_qp)
{
    int d0 = dct[0] + dct[1];
    int d1 = dct[2] + dct[3];
    int d2 = dct[0] - dct[1];
    int d3 = dct[2] - dct[3];
    int dmf = dequant_mf[i_qp % 6][0] << i_qp / 6;
    dct[0] = (d0 + d1) * dmf >> 5;
    dct[1] = (d0 - d1) * dmf >> 5;
    dct[2] = (d2 + d3) * dmf >> 5;
    dct[3] = (d2 - d3) * dmf >> 5;
}

void dct4x4dc(dctcoef d[16])
{
    dctcoef tmp[16];

    for (int i = 0; i < 4; i++)
    {
        int s01 = d[i * 4 + 0] + d[i * 4 + 1];
        int d01 = d[i * 4 + 0] - d[i * 4 + 1];
        int s23 = d[i * 4 + 2] + d[i * 4 + 3];
        int d23 = d[i * 4 + 2] - d[i * 4 + 3];

        tmp[0 * 4 + i] = s01 + s23;
        tmp[1 * 4 + i] = s01 - s23;
        tmp[2 * 4 + i] = d01 - d23;
        tmp[3 * 4 + i] = d01 + d23;
    }

    for (int i = 0; i < 4; i++)
    {
        int s01 = tmp[i * 4 + 0] + tmp[i * 4 + 1];
        int d01 = tmp[i * 4 + 0] - tmp[i * 4 + 1];
        int s23 = tmp[i * 4 + 2] + tmp[i * 4 + 3];
        int d23 = tmp[i * 4 + 2] - tmp[i * 4 + 3];

        d[i * 4 + 0] = (s01 + s23 + 1) >> 1;
        d[i * 4 + 1] = (s01 - s23 + 1) >> 1;
        d[i * 4 + 2] = (d01 - d23 + 1) >> 1;
        d[i * 4 + 3] = (d01 + d23 + 1) >> 1;
    }
}

void idct4x4dc(dctcoef d[16])
{
    dctcoef tmp[16];

    for (int i = 0; i < 4; i++)
    {
        int s01 = d[i * 4 + 0] + d[i * 4 + 1];
        int d01 = d[i * 4 + 0] - d[i * 4 + 1];
        int s23 = d[i * 4 + 2] + d[i * 4 + 3];
        int d23 = d[i * 4 + 2] - d[i * 4 + 3];

        tmp[0 * 4 + i] = s01 + s23;
        tmp[1 * 4 + i] = s01 - s23;
        tmp[2 * 4 + i] = d01 - d23;
        tmp[3 * 4 + i] = d01 + d23;
    }

    for (int i = 0; i < 4; i++)
    {
        int s01 = tmp[i * 4 + 0] + tmp[i * 4 + 1];
        int d01 = tmp[i * 4 + 0] - tmp[i * 4 + 1];
        int s23 = tmp[i * 4 + 2] + tmp[i * 4 + 3];
        int d23 = tmp[i * 4 + 2] - tmp[i * 4 + 3];

        d[i * 4 + 0] = s01 + s23;
        d[i * 4 + 1] = s01 - s23;
        d[i * 4 + 2] = d01 - d23;
        d[i * 4 + 3] = d01 + d23;
    }
}

// dct反变换
static void add4x4_idct(pixel *p_dst, dctcoef dct[16])
{
    dctcoef d[16];
    dctcoef tmp[16];

    for (int i = 0; i < 4; i++)
    {
        int s02 = dct[0 * 4 + i] + dct[2 * 4 + i];
        int d02 = dct[0 * 4 + i] - dct[2 * 4 + i];
        int s13 = dct[1 * 4 + i] + (dct[3 * 4 + i] >> 1);
        int d13 = (dct[1 * 4 + i] >> 1) - dct[3 * 4 + i];

        tmp[i * 4 + 0] = s02 + s13;
        tmp[i * 4 + 1] = d02 + d13;
        tmp[i * 4 + 2] = d02 - d13;
        tmp[i * 4 + 3] = s02 - s13;
    }

    for (int i = 0; i < 4; i++)
    {
        int s02 = tmp[0 * 4 + i] + tmp[2 * 4 + i];
        int d02 = tmp[0 * 4 + i] - tmp[2 * 4 + i];
        int s13 = tmp[1 * 4 + i] + (tmp[3 * 4 + i] >> 1);
        int d13 = (tmp[1 * 4 + i] >> 1) - tmp[3 * 4 + i];

        d[0 * 4 + i] = (s02 + s13 + 32) >> 6;
        d[1 * 4 + i] = (d02 + d13 + 32) >> 6;
        d[2 * 4 + i] = (d02 - d13 + 32) >> 6;
        d[3 * 4 + i] = (s02 - s13 + 32) >> 6;
    }

    for (int y = 0; y < 4; y++)
    {
        for (int x = 0; x < 4; x++)
            p_dst[x] = x264_clip_pixel(p_dst[x] + d[y * 4 + x]);
        p_dst += FDEC_STRIDE;
    }
}

void add8x8_idct(pixel *p_dst, dctcoef dct[4][16])
{
    add4x4_idct(&p_dst[0], dct[0]);
    add4x4_idct(&p_dst[4], dct[1]);
    add4x4_idct(&p_dst[4 * FDEC_STRIDE + 0], dct[2]);
    add4x4_idct(&p_dst[4 * FDEC_STRIDE + 4], dct[3]);
}

void add16x16_idct(pixel *p_dst, dctcoef dct[16][16])
{
    add8x8_idct(&p_dst[0], &dct[0]);
    add8x8_idct(&p_dst[8], &dct[4]);
    add8x8_idct(&p_dst[8 * FDEC_STRIDE + 0], &dct[8]);
    add8x8_idct(&p_dst[8 * FDEC_STRIDE + 8], &dct[12]);
}

static inline void add4x4_idct_dc(pixel *p_dst, dctcoef dc)
{
    dc = (dc + 32) >> 6;
    for (int i = 0; i < 4; i++, p_dst += FDEC_STRIDE)
    {
        p_dst[0] = x264_clip_pixel(p_dst[0] + dc);
        p_dst[1] = x264_clip_pixel(p_dst[1] + dc);
        p_dst[2] = x264_clip_pixel(p_dst[2] + dc);
        p_dst[3] = x264_clip_pixel(p_dst[3] + dc);
    }
}

void add8x8_idct_dc(pixel *p_dst, dctcoef dct[4])
{
    add4x4_idct_dc(&p_dst[0], dct[0]);
    add4x4_idct_dc(&p_dst[4], dct[1]);
    add4x4_idct_dc(&p_dst[4 * FDEC_STRIDE + 0], dct[2]);
    add4x4_idct_dc(&p_dst[4 * FDEC_STRIDE + 4], dct[3]);
}

void add16x16_idct_dc(pixel *p_dst, dctcoef dct[16])
{
    for (int i = 0; i < 4; i++, dct += 4, p_dst += 4 * FDEC_STRIDE)
    {
        add4x4_idct_dc(&p_dst[0], dct[0]);
        add4x4_idct_dc(&p_dst[4], dct[1]);
        add4x4_idct_dc(&p_dst[8], dct[2]);
        add4x4_idct_dc(&p_dst[12], dct[3]);
    }
}

void dct2x2dc(dctcoef d[4], dctcoef dct4x4[4][16])
{
    int d0 = dct4x4[0][0] + dct4x4[1][0];
    int d1 = dct4x4[2][0] + dct4x4[3][0];
    int d2 = dct4x4[0][0] - dct4x4[1][0];
    int d3 = dct4x4[2][0] - dct4x4[3][0];
    d[0] = d0 + d1;
    d[2] = d2 + d3;
    d[1] = d0 - d1;
    d[3] = d2 - d3;
    dct4x4[0][0] = 0;
    dct4x4[1][0] = 0;
    dct4x4[2][0] = 0;
    dct4x4[3][0] = 0;
}

void dct2x4dc(dctcoef dct[8], dctcoef dct4x4[8][16])
{
    int a0 = dct4x4[0][0] + dct4x4[1][0];
    int a1 = dct4x4[2][0] + dct4x4[3][0];
    int a2 = dct4x4[4][0] + dct4x4[5][0];
    int a3 = dct4x4[6][0] + dct4x4[7][0];
    int a4 = dct4x4[0][0] - dct4x4[1][0];
    int a5 = dct4x4[2][0] - dct4x4[3][0];
    int a6 = dct4x4[4][0] - dct4x4[5][0];
    int a7 = dct4x4[6][0] - dct4x4[7][0];
    int b0 = a0 + a1;
    int b1 = a2 + a3;
    int b2 = a4 + a5;
    int b3 = a6 + a7;
    int b4 = a0 - a1;
    int b5 = a2 - a3;
    int b6 = a4 - a5;
    int b7 = a6 - a7;
    dct[0] = b0 + b1;
    dct[1] = b2 + b3;
    dct[2] = b0 - b1;
    dct[3] = b2 - b3;
    dct[4] = b4 - b5;
    dct[5] = b6 - b7;
    dct[6] = b4 + b5;
    dct[7] = b6 + b7;
    dct4x4[0][0] = 0;
    dct4x4[1][0] = 0;
    dct4x4[2][0] = 0;
    dct4x4[3][0] = 0;
    dct4x4[4][0] = 0;
    dct4x4[5][0] = 0;
    dct4x4[6][0] = 0;
    dct4x4[7][0] = 0;
}
